#include <sys/stat.h>
#include <fstream>
#include <glog/logging.h>

#include "sample_detector.h"

static float BoxIOU(const cv::Rect2d &b1, const cv::Rect2d &b2)
{
    cv::Rect2d inter = b1 & b2;
    return inter.area() / (b1.area() + b2.area() - inter.area());
}


SampleDetector::SampleDetector()
{
}

bool SampleDetector::Init(float thresh, const std::string &model_name)
{
 
    m_net = migraphx::parse_onnx(model_name);
    std::unordered_map<std::string, migraphx::shape> inputs  = m_net.get_inputs();
    std::unordered_map<std::string, migraphx::shape> outputs = m_net.get_outputs();
    m_inputName                                                = inputs.begin()->first;
    m_inputShape                                               = inputs.begin()->second;
    SDKLOG(INFO)<<"inputShape:" << m_inputShape.lens().size();
    int N                                                    = m_inputShape.lens()[0];
    int C                                                    = m_inputShape.lens()[1];
    int H                                                    = m_inputShape.lens()[2];
    int W                                                    = m_inputShape.lens()[3];
    SDKLOG(INFO)<<"model input size(W, H):"<<W<<","<<H;
    m_inputSize                                                = cv::Size(W, H);
    m_outputName                                               = outputs.begin()->first;
    m_outputShape                                              = outputs.begin()->second;
    SDKLOG(INFO)<<"outputShape:"<< m_outputShape.lens().size();

    migraphx::target gpuTarget = migraphx::gpu::target{};
    migraphx::quantize_fp16(m_net);
    migraphx::compile_options options;
    options.device_id    = 0; // 设置GPU设备，默认为0号设备
    options.offload_copy = true;
    SDKLOG(INFO)<<"compiling ...";
    m_net.compile(gpuTarget, options);

    std::unordered_map<std::string, migraphx::argument> inputData;
    inputData[m_inputName] = migraphx::argument{m_inputShape};
    SDKLOG(INFO)<<"warming up ...";
    m_net.eval(inputData);


    SDKLOG(INFO)<<"init model done:"<<model_name;

    mThresh = thresh;
    return true;
}

bool SampleDetector::UnInit()
{
    return true;
}

SampleDetector::~SampleDetector()
{
    UnInit();
}

bool SampleDetector::ProcessImage(const cv::Mat &img, std::vector<BoxInfo> &det_results)
{
    cv::Mat in_mat = img.clone();
    float r = std::min(m_inputSize.width / static_cast<float>(in_mat.rows), m_inputSize.height / static_cast<float>(in_mat.cols));
    cv::Size new_size = cv::Size{static_cast<int>(in_mat.cols * r), static_cast<int>(in_mat.rows * r)};
    cv::Mat resized_mat;
    cv::resize(in_mat, resized_mat, new_size);
    if (m_board.empty())
    {
        m_board = cv::Mat(m_inputSize, CV_8UC3, cv::Scalar(114, 114, 114));
    }
    resized_mat.copyTo(m_board(cv::Rect{0, 0, resized_mat.cols, resized_mat.rows}));
    cv::cvtColor(m_board, m_board, cv::COLOR_BGR2RGB);
    m_board.convertTo(m_normalized_mat, CV_32FC3, 1 / 255.);

    cv::dnn::blobFromImage(m_normalized_mat, m_normalized_mat, 1.0, m_inputSize, cv::Scalar(0, 0, 0), false, false);

std::unordered_map<std::string, migraphx::argument> inputData;
    inputData[m_inputName] = migraphx::argument{m_inputShape, (float*)m_normalized_mat.data};

    std::vector<migraphx::argument> inferenceResults = m_net.eval(inputData);
    migraphx::argument infer_result = inferenceResults[0];

    float scale = std::min(m_inputSize.width / (in_mat.cols * 1.0), m_inputSize.height / (in_mat.rows * 1.0));
    std::vector<BoxInfo> detBoxes{};
    float *obuf = (float*)(infer_result.data());

    decode_outputs(obuf, detBoxes, scale, in_mat.cols, in_mat.rows);
    det_results.clear();
    det_results = BatchNMS(detBoxes);

    return true;
}

void SampleDetector::runNms(std::vector<BoxInfo> &vecBoxObjs)
{
    std::sort(vecBoxObjs.begin(), vecBoxObjs.end(), [](const BoxInfo &b1, const BoxInfo &b2)
              { return b1.score > b2.score; });
    for (int i = 0; i < vecBoxObjs.size(); ++i)
    {
        if (vecBoxObjs[i].score == 0)
        {
            continue;
        }
        for (int j = i + 1; j < vecBoxObjs.size(); ++j)
        {
            if (vecBoxObjs[j].score == 0)
            {
                continue;
            }
            cv::Rect pos1{vecBoxObjs[i].x1, vecBoxObjs[i].y1, vecBoxObjs[i].x2 - vecBoxObjs[i].x1, vecBoxObjs[i].y2 - vecBoxObjs[i].y1};
            cv::Rect pos2{vecBoxObjs[j].x1, vecBoxObjs[j].y1, vecBoxObjs[j].x2 - vecBoxObjs[j].x1, vecBoxObjs[j].y2 - vecBoxObjs[j].y1};
            if (BoxIOU(pos1, pos2) >= 0.45)
            {
                vecBoxObjs[j].score = 0;
            }
        }
    }
    for (auto iter = vecBoxObjs.begin(); iter != vecBoxObjs.end(); ++iter)
    {
        if (iter->score < 0.01)
        {
            vecBoxObjs.erase(iter);
            --iter;
        }
    }
}


std::vector<BoxInfo> SampleDetector::BatchNMS(std::vector<BoxInfo> &vecBoxObjs)
{
    std::vector<BoxInfo> singleClassBoxObjs;
    std::vector<BoxInfo> result;
    for (size_t classId = 0; classId < m_outputShape.lens()[2] - 5; classId++)
    {
        singleClassBoxObjs.clear();
        for (auto &obj : vecBoxObjs)
        {

            if (obj.label == int(classId))
            {
                singleClassBoxObjs.emplace_back(obj);
            }
        }
        runNms(singleClassBoxObjs);
        result.insert(result.end(), singleClassBoxObjs.begin(), singleClassBoxObjs.end());
    }
    return result;
}

void SampleDetector::decode_outputs(float *buffer,
                                    std::vector<BoxInfo> &objects,
                                    float scale,
                                    const int img_w,
                                    const int img_h)
{

    int box_num = m_outputShape.lens()[1];
    int class_num = m_outputShape.lens()[2] - 5;

    for (int i = 0; i < box_num; ++i)
    {
        int index = i * (class_num + 5);

        if (buffer[index + 4] > mThresh)
        {
            float x = buffer[index] / scale;
            float y = buffer[index + 1] / scale;
            float w = buffer[index + 2] / scale;
            float h = buffer[index + 3] / scale;
            float *max_cls_pos = std::max_element(buffer + index + 5, buffer + index + 5 + class_num);

            if ((*max_cls_pos) * buffer[index + 4] > mThresh)
            {
                cv::Rect box{(int)x -(int)(w / 2), int(y - h / 2), (int)w, (int)h};
                box = box & cv::Rect(0, 0, img_w - 1, img_h - 1);
                if (box.area() > 0)
                {
                    BoxInfo box_info = {float(box.x), (float)(box.y), (float)(box.x + box.width), (float)(box.y + box.height),
                    (*max_cls_pos) * buffer[index + 4], int(max_cls_pos - (buffer + index + 5))};
                    objects.push_back(box_info);
                }
            }
        }
    }
}

